# -*- coding: utf-8 -*-
# @Time : 2021-11-18 16:42
# @Author : lwb
# @Site : 
# @File : Transformer_process.py
import math
import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_sequence
class PositionEncoding(nn.Module):
    def __init__(self,d_model,dropout=0.1,max_len=512):
        """
        位置编码器类的初始化函数

        共有三个参数，分别是
        d_model：词嵌入维度
        dropout: dropout触发比率
        max_len：每个句子的最大长度
        """
        super(PositionEncoding, self).__init__()
        pe=torch.zeros(max_len,d_model)     #
        position=torch.arange(0,max_len,dtype=torch.float).unsqueeze(1)   # 扩充一个维度  即为 (max_len,1)
        div_term=torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))
        pe[:,0::2]=torch.sin(position*div_term)
        self.dropout=nn.Dropout(dropout)
        pe[:,1::2]=torch.cos(position*div_term)
        pe=pe.unsqueeze(0).transpose(0,1)                                # 在位置扩充一个维度，并且将0,1维度进行交换
        self.register_buffer('pe',pe)                                     # 不对位置编码层邱梯度
    def forward(self,x):
        x=x+self.pe[:x.size(0),:]           # 输入的词向量与位置编码相加
        return self.dropout(x)

def lenght_to_mask(lengths):
    """
    将序列长度转换成Mask矩阵
    print(lenght_to_mask(torch.tensor([3,5,4])))
    tensor([[ True,  True,  True, False, False],
        [ True,  True,  True,  True,  True],
        [ True,  True,  True,  True, False]])

    param lenghts: [batch,]
    return batch*max_len
    """
    max_len=torch.max(lengths)
    mask=torch.arange(max_len).expand(lengths.shape[0],max_len)<lengths.unsqueeze(1)
    return mask


# 这个函数也可不用，只需要将他的操作，在dataset中return 数据的时候应该就行了
# collate_fn参数指向一个函数，用于对一个批次的样本进行整理，如将其转换成张量
def collate_fn(examples):
    lenghts=torch.tensor([len(ex[0]) for ex in examples])
    inputs = [torch.tensor(ex[0]) for ex in examples]
    # 输出的目标targets为该批次中全部样例输出结果（0或者1）构成的张量
    targets = torch.tensor([ex[1] for ex in examples], dtype=torch.long)
    inputs = pad_sequence(inputs,batch_first=True)
    return inputs,lenghts, targets